package com.jd.genie.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jd.genie.agent.enums.AutoBotsResultStatus;
import com.jd.genie.agent.enums.ResponseTypeEnum;
import com.jd.genie.config.GenieConfig;
import com.jd.genie.handler.AgentHandlerFactory;
import com.jd.genie.handler.AgentHandlerService;
import com.jd.genie.model.dto.AutoBotsResult;
import com.jd.genie.model.multi.EventResult;
import com.jd.genie.model.req.AgentRequest;
import com.jd.genie.model.req.GptQueryReq;
import com.jd.genie.model.response.AgentResponse;
import com.jd.genie.model.response.GptProcessResult;
import com.jd.genie.util.ChateiUtils;
import com.jd.genie.util.SseUtil;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.concurrent.TimeUnit;

@Slf4j
@Service
public class IGptProcessService {

    @Autowired
    private GenieConfig genieConfig;

    @Resource
    private AgentHandlerFactory agentHandlerFactory;

    /**
     * 单智能体，多智能体 Agent 增量接口.
     */
    public SseEmitter queryMultiAgentIncrStream(GptQueryReq req)
    {
        long timeoutMillis = TimeUnit.HOURS.toMillis(1);
        req.setUser("genie");
        String traceId = ChateiUtils.getRequestId(req);
        req.setTraceId(traceId);
        final SseEmitter emitter = SseUtil.build(timeoutMillis, req.getTraceId());
        searchForAgentRequest(req, emitter);
        log.info("queryMultiAgentIncrStream GptQueryReq request:{}", req);
        return emitter;
    }

    /**
     * 请求多 agent发送请求入口函数.
     */
    public AutoBotsResult searchForAgentRequest(GptQueryReq gptQueryReq, SseEmitter sseEmitter)
    {
        AgentRequest agentRequest = buildAgentRequest(gptQueryReq);
        log.info("{} start handle Agent request: {}", gptQueryReq.getRequestId(), JSON.toJSONString(agentRequest));
        try {
            handleMultiAgentRequest(agentRequest, sseEmitter);
        } catch (Exception e) {
            log.error("{}, error in requestMultiAgent, deepThink: {}, errorMsg: {}", gptQueryReq.getRequestId(), gptQueryReq.getDeepThink(), e.getMessage(), e);
            throw e;
        } finally {
            log.info("{}, agent.query.web.singleRequest end, requestId: {}", gptQueryReq.getRequestId(), JSON.toJSONString(gptQueryReq));
        }

        return ChateiUtils.toAutoBotsResult(agentRequest, AutoBotsResultStatus.loading.name());
    }

    @SneakyThrows
    public void handleMultiAgentRequest(AgentRequest autoReq, SseEmitter sseEmitter)
    {
        long startTime = System.currentTimeMillis();

        Request request;
        {
            String reqId = autoReq.getRequestId();
            autoReq.setRequestId(autoReq.getRequestId());
            String url = "http://127.0.0.1:8080/AutoAgent";
            RequestBody body = RequestBody.create(
                    MediaType.parse("application/json"),
                    JSONObject.toJSONString(autoReq)
            );
            autoReq.setRequestId(reqId);
            request = new Request.Builder().url(url).post(body).build();
        }

        // todo lq
        // genieControllerInternal.AutoAgent(autoReq);

        log.info("{} agentRequest:{}", autoReq.getRequestId(), JSON.toJSONString(request));
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(60, TimeUnit.SECONDS) // 设置连接超时时间为 60 秒
                .readTimeout(genieConfig.getSseClientReadTimeout(), TimeUnit.SECONDS)    // 设置读取超时时间为 60 秒
                .writeTimeout(1800, TimeUnit.SECONDS)   // 设置写入超时时间为 60 秒
                .callTimeout(genieConfig.getSseClientConnectTimeout(), TimeUnit.SECONDS)    // 设置调用超时时间为 60 秒
                .build();

        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            if (responseBody == null) {
                log.error("{} auto agent empty response body", autoReq.getRequestId());
                return;
            }
            if (!response.isSuccessful()) {
                log.error("{}, response body is failed: {}", autoReq.getRequestId(), responseBody.string());
                return;
            }

            String line;
            BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()));
            while ((line = reader.readLine()) != null) {
                if (!line.startsWith("data:")) {
                    continue;
                }

                String data = line.substring(5);
                if (data.equals("[DONE]")) {
                    log.info("{} data equals with [DONE] {}:", autoReq.getRequestId(), data);
                    break;
                }

                if (data.startsWith("heartbeat")) {
                    GptProcessResult result = buildHeartbeatData(autoReq.getRequestId());
                    sseEmitter.send(result);
                    log.info("{} heartbeat-data: {}", autoReq.getRequestId(), data);
                    continue;
                }

                log.info("{} recv from autocontroller: {}", autoReq.getRequestId(), data);
                AgentResponse agentResponse = JSON.parseObject(data, AgentResponse.class);

                AgentHandlerService handler = agentHandlerFactory.getHandler(autoReq.getAgentType());
                GptProcessResult result = handler.handleResponse(autoReq, agentResponse, new ArrayList<>(), new EventResult());
                sseEmitter.send(result);
                if (result.isFinished()) {
                    // 记录任务执行时间
                    log.info("task complete {} task total cost time:{}ms", autoReq.getRequestId(), System.currentTimeMillis() - startTime);
                    sseEmitter.complete();
                }
            }
        } catch (Exception e) {
            log.error("", e);
        }
    }

    private AgentRequest buildAgentRequest(GptQueryReq req)
    {
        AgentRequest request = new AgentRequest();
        request.setRequestId(req.getTraceId());
        request.setErp(req.getUser());
        request.setQuery(req.getQuery());
        request.setAgentType(req.getDeepThink() == 0 ? 5 : 3);
        request.setIsStream(true);
        request.setOutputStyle(req.getOutputStyle());
        return request;
    }

    private GptProcessResult buildHeartbeatData(String requestId)
    {
        GptProcessResult result = new GptProcessResult();
        result.setFinished(false);
        result.setStatus("success");
        result.setResponseType(ResponseTypeEnum.text.name());
        result.setResponse("");
        result.setResponseAll("");
        result.setUseTimes(0);
        result.setUseTokens(0);
        result.setReqId(requestId);
        result.setPackageType("heartbeat");
        result.setEncrypted(false);
        return result;
    }

}
